# -*- coding: UTF-8 -*-
'''
@Project ：automation_-interface_testing
@File    ：mysql_operate.py
@IDE     ：PyCharm 
@Author  ：思亦忆
@Date    ：2025/9/5 星期五 21:58:06

'''
"""
读取config.yaml文件，获取mysql配置项。
封装一个数据库类，实现连接、断开、执行SQL等方法。
使用yaml_config.py（假设存在）来读取配置，或者我们直接在数据库模块中读取配置。
"""

import pymysql
from pymysql import cursors
from pymysql.err import OperationalError, MySQLError
import yaml
from pathlib import Path
import logging
from src.common.logs import get_logger

# 获取日志记录器
logger = get_logger()


class MySQLOperator:
    """MySQL数据库操作类"""

    def __init__(self, config_path='config/config.yaml'):
        """
        初始化数据库连接

        Args:
            config_path: 配置文件路径，默认为'config/config.yaml'
        """
        self.config_path = config_path
        self.connection = None
        self.cursor = None
        self.db_config = self._load_db_config()

    def _load_db_config(self):
        """
        从YAML配置文件加载数据库配置

        Returns:
            dict: 数据库配置字典
        """
        try:
            # 构建配置文件路径
            config_file = Path(self.config_path)
            if not config_file.exists():
                # 尝试从项目根目录查找
                project_root = Path(__file__).parent.parent
                config_file = project_root / self.config_path
                if not config_file.exists():
                    raise FileNotFoundError(f"数据库配置文件不存在: {self.config_path}")

            # 读取配置文件
            with open(config_file, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            # 提取数据库配置
            db_config = config.get('mysql', {})

            # 验证必要配置项
            required_keys = ['host', 'port', 'user', 'password', 'db']
            missing_keys = [key for key in required_keys if key not in db_config]

            if missing_keys:
                raise ValueError(f"数据库配置缺少必要参数: {missing_keys}")

            logger.info(f"成功加载数据库配置: {db_config['host']}:{db_config['port']}/{db_config['db']}")
            return db_config

        except Exception as e:
            logger.error(f"加载数据库配置失败: {e}")
            raise

    def connect(self):
        """
        连接到MySQL数据库

        Returns:
            bool: 连接是否成功
        """
        try:
            # 确保密码是字符串类型
            password = str(self.db_config.get('password', ''))

            # 建立数据库连接
            self.connection = pymysql.connect(
                host=self.db_config['host'],
                port=self.db_config['port'],
                user=self.db_config['user'],
                password=password,  # 确保密码是字符串
                database=self.db_config['db'],
                charset='utf8mb4',
                cursorclass=cursors.DictCursor,
                autocommit=True
            )

            # 创建游标
            self.cursor = self.connection.cursor()

            logger.info("数据库连接成功")
            return True

        except OperationalError as e:
            logger.error(f"数据库连接失败: {e}")
            return False
        except Exception as e:
            logger.error(f"数据库连接异常: {e}")
            return False

    def disconnect(self):
        """断开数据库连接"""
        try:
            if self.cursor:
                self.cursor.close()
            if self.connection:
                self.connection.close()
            logger.info("数据库连接已关闭")
        except Exception as e:
            logger.error(f"关闭数据库连接时出错: {e}")

    def execute_query(self, sql, params=None):
        """
        执行查询语句

        Args:
            sql: SQL查询语句
            params: 查询参数

        Returns:
            list: 查询结果列表
        """
        try:
            # 确保数据库连接
            if not self.connection or not self.connection.open:
                self.connect()

            # 执行查询
            self.cursor.execute(sql, params)
            result = self.cursor.fetchall()

            logger.debug(f"执行查询: {sql}, 参数: {params}, 结果行数: {len(result)}")
            return result

        except MySQLError as e:
            logger.error(f"查询执行失败: {e}, SQL: {sql}, 参数: {params}")
            raise
        except Exception as e:
            logger.error(f"查询执行异常: {e}")
            raise

    def execute_non_query(self, sql, params=None):
        """
        执行非查询语句（INSERT, UPDATE, DELETE）

        Args:
            sql: SQL语句
            params: 参数

        Returns:
            int: 受影响的行数
        """
        try:
            # 确保数据库连接
            if not self.connection or not self.connection.open:
                self.connect()

            # 执行非查询语句
            affected_rows = self.cursor.execute(sql, params)

            logger.debug(f"执行非查询: {sql}, 参数: {params}, 影响行数: {affected_rows}")
            return affected_rows

        except MySQLError as e:
            logger.error(f"非查询执行失败: {e}, SQL: {sql}, 参数: {params}")
            raise
        except Exception as e:
            logger.error(f"非查询执行异常: {e}")
            raise

    def execute_many(self, sql, params_list):
        """
        批量执行SQL语句

        Args:
            sql: SQL语句
            params_list: 参数列表

        Returns:
            int: 受影响的总行数
        """
        try:
            # 确保数据库连接
            if not self.connection or not self.connection.open:
                self.connect()

            # 批量执行
            affected_rows = self.cursor.executemany(sql, params_list)

            logger.debug(f"批量执行: {sql}, 参数数量: {len(params_list)}, 影响行数: {affected_rows}")
            return affected_rows

        except MySQLError as e:
            logger.error(f"批量执行失败: {e}, SQL: {sql}")
            raise
        except Exception as e:
            logger.error(f"批量执行异常: {e}")
            raise

    def get_table_columns(self, table_name):
        """
        获取表的列信息

        Args:
            table_name: 表名

        Returns:
            list: 列信息列表
        """
        sql = f"DESCRIBE {table_name}"
        return self.execute_query(sql)

    def get_table_data(self, table_name, condition=None, limit=None):
        """
        获取表数据

        Args:
            table_name: 表名
            condition: 查询条件
            limit: 限制返回行数

        Returns:
            list: 数据列表
        """
        sql = f"SELECT * FROM {table_name}"
        if condition:
            sql += f" WHERE {condition}"
        if limit:
            sql += f" LIMIT {limit}"

        return self.execute_query(sql)

    def insert_data(self, table_name, data):
        """
        插入数据

        Args:
            table_name: 表名
            data: 数据字典

        Returns:
            int: 受影响的行数
        """
        columns = ', '.join(data.keys())
        placeholders = ', '.join(['%s'] * len(data))
        sql = f"INSERT INTO {table_name} ({columns}) VALUES ({placeholders})"

        return self.execute_non_query(sql, list(data.values()))

    def update_data(self, table_name, data, condition):
        """
        更新数据

        Args:
            table_name: 表名
            data: 要更新的数据字典
            condition: 更新条件

        Returns:
            int: 受影响的行数
        """
        set_clause = ', '.join([f"{key} = %s" for key in data.keys()])
        sql = f"UPDATE {table_name} SET {set_clause} WHERE {condition}"

        return self.execute_non_query(sql, list(data.values()))

    def delete_data(self, table_name, condition):
        """
        删除数据

        Args:
            table_name: 表名
            condition: 删除条件

        Returns:
            int: 受影响的行数
        """
        sql = f"DELETE FROM {table_name} WHERE {condition}"
        return self.execute_non_query(sql)

    def __enter__(self):
        """上下文管理器入口"""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出"""
        self.disconnect()


# 单例模式获取数据库操作实例
_db_instance = None


def get_db_instance(config_path='config/config.yaml'):
    """获取数据库操作实例（单例模式）"""
    global _db_instance

    if _db_instance is None:
        _db_instance = MySQLOperator(config_path)
        _db_instance.connect()

    return _db_instance


# 示例用法
if __name__ == "__main__":
    # 使用上下文管理器（推荐）
    with MySQLOperator() as db:
        # 1. 查看当前数据库
        current_db = db.execute_query("SELECT DATABASE() as current_db")
        print("当前数据库:", current_db)

        # 2. 查看所有数据库
        databases = db.execute_query("SHOW DATABASES")
        print("所有数据库:", [d['Database'] for d in databases])

        # 3. 查看当前数据库中的所有表
        tables = db.execute_query("SHOW TABLES")
        print("当前数据库中的表:", tables)

        # 4. 如果有表，查询第一张表的数据
        if tables:
            # 获取第一张表的表名
            table_name = list(tables[0].values())[0]
            result = db.execute_query(f"SELECT * FROM `{table_name}` LIMIT 5")
            print(f"表 {table_name} 的数据:", result)
